"""DeepAgent实现"""
import os
from typing import Literal
from tavily import TavilyClient
from deepagents import create_deep_agent
from langchain_core.messages import BaseMessage
from config.config import OPENAI_API_KEY, OPENAI_BASE_URL, DEFAULT_MODEL
from datetime import datetime, timedelta, timezone


def get_beijing_time_str() -> tuple[str, str]:
    """获取北京时间和星期"""
    beijing_offset = timezone(timedelta(hours=8))
    beijing_time = datetime.now(beijing_offset)
    time_str = beijing_time.strftime("%Y年%m月%d日 %H:%M:%S")
    weekday = ["周一", "周二", "周三", "周四", "周五", "周六", "周日"][beijing_time.weekday()]
    return time_str, weekday


def inject_time_info(base_prompt: str) -> str:
    """注入时间信息到提示词"""
    time_str, weekday = get_beijing_time_str()
    return f"""【当前北京时间】{time_str} {weekday}

{base_prompt}"""


def internet_search(
    query: str,
    max_results: int = 5,
    topic: Literal["general", "news", "finance"] = "general",
    include_raw_content: bool = False,
):
    """Run a web search"""
    tavily_api_key = os.environ.get("TAVILY_API_KEY")
    if not tavily_api_key:
        return "错误: 未配置 TAVILY_API_KEY"
    
    tavily_client = TavilyClient(api_key=tavily_api_key)
    return tavily_client.search(
        query,
        max_results=max_results,
        include_raw_content=include_raw_content,
        topic=topic,
    )


def create_deep_agent_instance(
    system_prompt: str = None,
    model_name: str = None,
    api_key: str = None,
    base_url: str = None
):
    """创建DeepAgent实例
    
    Args:
        system_prompt: 系统提示词
        model_name: 模型名称
        api_key: API密钥
        base_url: API基础URL
        
    Returns:
        配置好的DeepAgent
    """
    if model_name is None:
        model_name = DEFAULT_MODEL
    
    if system_prompt is None:
        system_prompt = """You are an expert researcher. Your job is to conduct thorough research and then write a polished report.

You have access to an internet search tool as your primary means of gathering information.

## `internet_search` 

Use this to run an internet search for a given query. You can specify the max number of results to return, the topic, and whether raw content should be included."""
    
    system_prompt = inject_time_info(system_prompt)
    
    llm_kwargs = {
        "model": model_name,
        "temperature": 0
    }
    
    if api_key:
        llm_kwargs["api_key"] = api_key
    elif OPENAI_API_KEY:
        llm_kwargs["api_key"] = OPENAI_API_KEY
    
    if base_url:
        llm_kwargs["base_url"] = base_url
    elif OPENAI_BASE_URL:
        llm_kwargs["base_url"] = OPENAI_BASE_URL
    
    os.environ["OPENAI_API_KEY"] = llm_kwargs["api_key"]
    if "base_url" in llm_kwargs:
        os.environ["OPENAI_BASE_URL"] = llm_kwargs["base_url"]
    
    from langchain_openai import ChatOpenAI
    llm = ChatOpenAI(**llm_kwargs)
    
    agent = create_deep_agent(
        tools=[internet_search],
        system_prompt=system_prompt,
        model=llm
    )
    
    return agent
